import pandas as pd
from matplotlib import pyplot as plt
from algo import *

class Strategy(Model):
    """
    run()计算样本多因子模型策略表现
    超额收益净值曲线储存为Portfolio Performance.png
    """
    selected_ft = {
        # 估值
        'bp_lf_plus': -1,
        # 质量
        'roe_qr_plus_r': 1,
        'eps_qr_plus_r': 1,
        # 成长
        'netprofit_qr_yoy_plus_r': 1,
        # 分析师
        'sue_minus_con_qr_plus': 1,
    }
    lag = 1
    grp = 10  # 分组数

    def __init__(self):
        super().__init__()
        self.read_idx_eod()
        self.calc_return()
        self.ft_test: (None, pd.DataFrame) = None

    def run(self):
        # 样本外多因子选股
        self.portf = self.generate_portf()
        # 组合收益
        self.portf_performance(self.portf)


    def generate_portf(self):
        # 计算因子值
        ft_list = list(Strategy.selected_ft.keys())
        uni_test = self.uni[self.uni['date'] > self.train_test_split].copy(deep=False)
        self.ft_test = self.calc_ft(uni_test, ft_list)
        self.ft_test = self.ft_test.loc[:,['date', 'code']+ft_list]
        direction = list(Strategy.selected_ft.values())
        self.ft_test.loc[:, ft_list] = self.ft_test[ft_list] * direction

        # 多因子正交化处理
        def orth(x:pd.DataFrame):
            for i in range(1, len(ft_list)):
                x[ft_list[i]] = FTProcessor.orthogonal(x[ft_list[i-1]].values, x[ft_list[i]].values)
            return x
        self.ft_test = self.ft_test.groupby('date').apply(lambda x: orth(x))

        # 等权打分
        score = self.ft_test.set_index(['date', 'code']).mean(axis=1).reset_index()
        score.columns = ['date', 'code', 'score']
        score.sort_values(by=['date', 'score'], inplace=True)

        # 多头组合
        def inner(x: pd.DataFrame):
            num_codes = int(np.round(x.shape[0] / self.grp, 0))
            assert num_codes > 0
            x = x.iloc[-num_codes:, :]
            return x

        portf = score.groupby('date', as_index=False).apply(lambda x: inner(x)).reset_index(drop=True)
        portf = portf[portf['date'] < portf['date'].max()]

        return portf

    def portf_performance(self, portf):
        # 匹配收益率日期
        map = {}
        for i, v in enumerate(self.dates[:-self.lag]):
            map[v] = self.dates[self.dates.index(v) + self.lag]
        portf.replace({'date': map}, inplace=True)

        # 合并收益率
        col_name = 'pct_chg_vwap'
        portf = portf.merge(self.uni[['date', 'code', col_name]].copy(deep=False), how='left', on=['date', 'code'])

        # 多头收益
        ret_portf = portf.dropna()[['date', col_name]].groupby('date').mean()
        # 超额收益
        ret_excess = ret_portf.merge(self.idx, on='date', how='left')
        ret_excess.iloc[:, 0] = ret_excess.iloc[:, 0] - ret_excess.iloc[:, 1]
        ret_excess = ret_excess.iloc[:, :1].dropna()

        # 画图
        nav = (1 + ret_excess).cumprod()
        plt.plot(nav)
        title = "Portfolio Performance"
        plt.title(title)
        plt.legend(['Excess Return'])
        plt.axhline(y=1, linestyle=':', color='grey')
        plt.savefig(f"{title}.png")
        plt.show(block=True)
        print('.')


if __name__ == '__main__':
    strategy = Strategy()
    strategy.run()






